package com.erebus.config;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.time.LocalDate;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY;

@Component
public class AiWebSocketHandler extends TextWebSocketHandler {

    @Autowired
    private ChatClient chatClient; // 假设你已经注入了 chatClient

    @Override
    public void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        String userMessage = message.getPayload();

        Flux<String> stream = chatClient.prompt()
                .user(userMessage)
                .system(s -> s.param("current_date", LocalDate.now().toString()))
                .advisors(a -> a.param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 18))
                .stream()
                .content();

        Flux<String> jsonStream = stream.map(chunk ->
                "{\"type\":\"text\",\"content\":\"" + escapeJson(chunk) + "\"}"
        );

        Flux<String> finalStream = jsonStream
                .startWith("{\"type\":\"start\"}")
                .concatWith(Flux.just("{\"type\":\"complete\"}"));

        finalStream.subscribe(
                chunk -> {
                    try {
                        session.sendMessage(new TextMessage(chunk));
                    } catch (IOException e) {
                        e.printStackTrace();
                    }
                },
                throwable -> {
                    try {
                        session.sendMessage(new TextMessage("{\"type\":\"error\",\"message\":\"" + throwable.getMessage() + "\"}"));
                    } catch (IOException e) {
                        e.printStackTrace();
                    }
                },
                () -> System.out.println("流完成")
        );
    }
    private String escapeJson(String str) {
        return str.replace("\\", "\\\\")
                .replace("\"", "\\\"")
                .replace("\n", "\\n")
                .replace("\r", "\\r");
    }
}